from tensorflow.keras import datasets, layers, models
import tensorflow as tf
import client
import server
import numpy as np

# 定义一个卷积神经网络（CNN）模型，用于训练和测试
def CNN_model_factory():
  model = models.Sequential()
  # 第一层卷积层，30 个卷积核，卷积核大小为 3x3，激活函数为 relu
  model.add(layers.Conv2D(30, (3, 3), activation='relu', input_shape=(28, 28, 1)))
  # 第一层池化层，池化窗口大小为 2x2
  model.add(layers.MaxPooling2D((2, 2)))
  # 第二层卷积层，50 个卷积核，卷积核大小为 3x3，激活函数为 relu
  model.add(layers.Conv2D(50, (3, 3), activation='relu'))
  # 第二层池化层，池化窗口大小为 2x2
  model.add(layers.MaxPooling2D((2, 2)))
  # 将多维数据展平为一维向量
  model.add(layers.Flatten())
  # 全连接层，100 个神经元，激活函数为 relu
  model.add(layers.Dense(100, activation='relu'))
  # 输出层，10 个神经元，激活函数为 softmax，用于输出分类概率
  model.add(layers.Dense(10,activation='softmax'))
  return model


#正常运行联邦学习。
def run_no_attacks(root_data, root_label, clients_data, clients_label, nclient, x_test, y_test):
  # 学习率相关超参数
  beta  = 5e-2
  alpha = 5e-2/beta
  # 模型工厂函数，用于创建 CNN 模型
  model_factory = CNN_model_factory
  # 存储客户端对象
  clients = []
  # 创建客户端对象
  for i in range(0,nclient):
    clients.append(client.Client(model_factory, clients_data[i], clients_label[i], learning_rate=beta, R=1, batch_size=32))
  # 创建根服务器对象
  root_client = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)
  # server初始化
  server1 = server.Server(model_factory, select_client=nclient, iteration=200, alpha=alpha, beta=beta)
  # 调用server的 train 方法进行训练
  server1.train(clients,root_client,'no_attacks',x_test, y_test)
  return server1

#仅用根数据集训练（基准对比）。
def run_only_server(root_data, root_label, x_test, y_test):
  # 定义学习率 beta
  beta  = 5e-2
  # 计算超参数 alpha
  alpha = 5e-2/beta
  # 获取模型工厂函数
  model_factory = CNN_model_factory
  # 创建服务器对象
  server1 = server.Server(model_factory, select_client=100, iteration=200, alpha=alpha, beta=beta)
  # 创建根客户端对象
  root_client = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)
  # 调用server的 train_self 方法仅在根数据集上进行训练
  server1.train_self(root_client,'train_on_server', x_test, y_test)
  return server1

def run_adaptive_attack(root_data, root_label, clients_data, clients_label, nclient, num_malicious, target_cosine_sim,x_test, y_test):
  """
  设置并运行带有自适应攻击的FLTrust实验。

  Args:
      root_data, root_label: 服务器根数据和标签。
      clients_data, clients_label: 所有客户端的数据和标签列表。
      nclient (int): 客户端总数。
      num_malicious (int): 恶意客户端的数量。
      x_test, y_test: 测试数据和标签。

  Returns:
      Server: 训练完成后的服务器对象。
  """
  # 设置学习率等超参数 (与无攻击场景保持一致)
  beta = 5e-2
  alpha = 5e-2 / beta
  model_factory = CNN_model_factory # 使用相同的模型结构

  # 创建所有客户端对象
  clients = []
  for i in range(nclient):
    # 客户端本身不需要知道自己是否是恶意的，因为攻击由服务器端的模拟器协调
    clients.append(client.Client(model_factory, clients_data[i], clients_label[i], learning_rate=beta, R=1, batch_size=32))

  # 创建根客户端对象
  root_client = client.Client(model_factory, root_data, root_label, learning_rate=beta, R=1, batch_size=32)

  # 创建服务器对象
  # 假设所有客户端都参与每轮训练 (select_client=nclient)
  # iteration 设置为与其他实验一致的轮数
  server_adaptive = server.Server(model_factory, select_client=nclient, iteration=200, alpha=alpha, beta=beta)

  # 修改 expr_basename 以包含 target_cosine_sim 值
  exp_id = f'adaptive_attack_{num_malicious}mal_cos{target_cosine_sim:.2f}'
  print(f"\n--- 开始运行FLTrust ({exp_id}) ---")

  server_adaptive.train_adaptive_attack(
      clients=clients,              # 传入所有客户端列表
      root_client=root_client,      # 传入根客户端
      num_malicious=num_malicious,  # 告知服务器有多少个恶意客户端
      expr_basename=exp_id, # 日志标识
    target_cosine_sim=target_cosine_sim,
      x_test=x_test,
      y_test=y_test
  )
  print(f"--- 完成FLTrust ({exp_id}) 的训练 ---")
  # 返回训练后的服务器对象
  return server_adaptive